# export_to_pyg.py 
import torch
import pandas as pd
from sklearn.preprocessing import StandardScaler
from torch_geometric.data import Data
from neo4j import GraphDatabase

# === Neo4j 配置 ===
URI = "bolt://localhost:7687"
USER = "neo4j"
PASSWORD = "fyx123"

driver = GraphDatabase.driver(URI, auth=(USER, PASSWORD))

def fetch_data():
    with driver.session() as session:
        # 获取所有节点（带属性）
        node_query = """
        MATCH (n)
        RETURN 
            id(n) AS node_id,
            labels(n)[0] AS label,
            n.id AS raw_id,
            n.role AS role,
            n.age AS age,
            n.violation_count AS violation_count,
            n.is_resigned_agent AS is_resigned_agent,
            n.has_bad_record AS has_bad_record,
            n.premium AS premium,
            n.amount AS amount,
            n.cost AS cost,
            n.hospital_days AS hospital_days
        """
        nodes = session.run(node_query)
        nodes_df = pd.DataFrame([dict(r) for r in nodes])

        # 获取所有边
        edge_query = "MATCH (s)-[r]->(t) RETURN id(s) AS src, id(t) AS dst"
        edges = session.run(edge_query)
        edges_df = pd.DataFrame([dict(r) for r in edges])

        return nodes_df, edges_df

def build_pyg_data(nodes_df, edges_df):
    # 节点ID到索引映射
    node_id_to_idx = {nid: idx for idx, nid in enumerate(nodes_df['node_id'])}
    nodes_df['idx'] = nodes_df['node_id'].map(node_id_to_idx)

    # 特征工程
    features = []

    # 角色 one-hot（简化：用 LabelEncoder）
    from sklearn.preprocessing import LabelEncoder
    role_encoder = LabelEncoder()
    nodes_df['role_encoded'] = role_encoder.fit_transform(nodes_df['role'].fillna('unknown'))
    features.append('role_encoded')

    # 数值特征
    num_features = ['age', 'violation_count', 'premium', 'amount', 'cost', 'hospital_days']
    for col in num_features:
        nodes_df[col] = pd.to_numeric(nodes_df[col], errors='coerce').fillna(0)
    features.extend(num_features)

    # 布尔特征
    bool_features = ['is_resigned_agent', 'has_bad_record']
    for col in bool_features:
        nodes_df[col] = nodes_df[col].fillna(False).astype(int)
    features.extend(bool_features)

    # 归一化数值特征
    scaler = StandardScaler()
    nodes_df[features] = scaler.fit_transform(nodes_df[features].values)

    # 构建特征矩阵
    X = nodes_df[features].values
    x = torch.tensor(X, dtype=torch.float)

    # 构建边索引
    edges_df['src_idx'] = edges_df['src'].map(node_id_to_idx)
    edges_df['dst_idx'] = edges_df['dst'].map(node_id_to_idx)
    edges_df = edges_df.dropna()
    edge_index = torch.tensor(edges_df[['src_idx', 'dst_idx']].values.T, dtype=torch.long)

    # 节点标签（Claim 节点的 is_fraud）
    y = torch.full((len(nodes_df),), -1, dtype=torch.long)
    # 注意：如果你的 Claim 节点没有 is_fraud 属性，先跳过标签
    try:
        fraud_query = "MATCH (c:Claim) RETURN id(c) AS node_id, c.is_fraud AS is_fraud"
        with driver.session() as session:
            fraud_nodes = session.run(fraud_query)
            fraud_map = {r["node_id"]: 1 if r["is_fraud"] else 0 for r in fraud_nodes}
        for i, row in nodes_df.iterrows():
            if row['label'] == 'Claim' and row['node_id'] in fraud_map:
                y[i] = fraud_map[row['node_id']]
    except:
        print("⚠️ Claim 节点无 is_fraud 属性，y 全为 -1")

    data = Data(x=x, edge_index=edge_index, y=y)
    torch.save(data, "fraud_kg_pyg_data.pt")
    torch.save(node_id_to_idx, "node_id_to_idx.pt")
    nodes_df.to_csv("nodes_for_pyg.csv", index=False)

    print(f"✅ PyG Dataset 已保存: {data}")
    return data

if __name__ == "__main__":
    print("🔄 从 Neo4j 导出数据...")
    nodes_df, edges_df = fetch_data()
    print(f"导出 {len(nodes_df)} 节点, {len(edges_df)} 边")
    build_pyg_data(nodes_df, edges_df)